import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np
from sklearn.utils import shuffle
from sklearn.model_selection import train_test_split
import time

# 加载MNIST数据集，这里先把所有数据合并到一起
mnist = tf.keras.datasets.mnist
(X_all, y_all), _ = mnist.load_data()

# 数据归一化，明确转换数据类型为float32，确保后续操作不出问题
X_all = X_all.astype(np.float32) / 255.0
# 增加通道维度，使用更明确的reshape方式
X_all = X_all.reshape(X_all.shape[0], X_all.shape[1], X_all.shape[2], 1)

# 不使用固定随机种子划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X_all, y_all, test_size = 0.2)

# 设置随机种子为当前时间（以确保每次运行数据打乱都不同）
np.random.seed(int(time.time()))
# 随机打乱训练数据（可选操作，取决于是否需要进一步打乱）
X_train, y_train = shuffle(X_train, y_train)

model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(X_train, y_train, epochs=5, validation_data=(X_test, y_test))

# 模型评估
loss, accuracy = model.evaluate(X_test, y_test)
print(f'Accuracy: {accuracy}')

# 预测
predictions = model.predict(X_test)

# 绘制前10个测试样本及其预测结果
for i in range(10):
    plt.subplot(2, 5, i + 1)
    plt.imshow(X_test[i].reshape(28, 28), cmap='gray')
    plt.title(f'Pred: {np.argmax(predictions[i])}\nTrue: {y_test[i]}')
    plt.axis('off')
plt.show()